{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch训练后动态量化BERT\n",
    " 本脚本展示使用HuggingFace的Bert模型进行训练后的动态量化, 代码参考如下连接: https://percent4.github.io/archives/\n",
    "\n",
    "                        \n",
    "                                                                            Author： Xian Yang\n",
    "                                                                            Date Time： 2023.11.27\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using the latest cached version of the module from /home/yangxianpku/.cache/huggingface/modules/datasets_modules/datasets/imdb/d613c88cf8fa3bab83b4ded3713f1f74830d1100e171db75bbddb80b3345c9c0 (last modified on Mon Nov 27 14:43:51 2023) since it couldn't be found locally at imdb., or remotely on the Hugging Face Hub.\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import torch\n",
    "import datasets\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "# 导入transformers库中的自动tokenizer工具，含padding功能的dataloader工具\n",
    "from transformers    import AutoTokenizer, DataCollatorWithPadding\n",
    "from transformers    import AutoModelForSequenceClassification\n",
    "from sklearn.metrics import accuracy_score, precision_recall_fscore_support\n",
    "\n",
    "MAX_LENGTH = 300\n",
    "\n",
    "device     = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "checkpoint = f\"../outputs/imdb_trainer/checkpoint-2346\"\n",
    "model      = AutoModelForSequenceClassification.from_pretrained(checkpoint).to(device)\n",
    "\n",
    "tokenizer  = AutoTokenizer.from_pretrained(checkpoint)\n",
    "\n",
    "# def tokenize_function(sample):\n",
    "#     return tokenizer(sample['text'], max_length=300, truncation=True)\n",
    "\n",
    "# 对数据集进行批量tokenize操作\n",
    "raw_datasets       = datasets.load_dataset('imdb')['test']     \n",
    "# data_collator      = DataCollatorWithPadding(tokenizer=tokenizer)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 量化前性能与精度评估\n",
    "\n",
    "### 1.1 量化前性能评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "avg time:  3.421308927536011\n"
     ]
    }
   ],
   "source": [
    "type(raw_datasets)\n",
    "\n",
    "true_labels, pred_labels = [], [] \n",
    "inputs = []\n",
    "for item in raw_datasets:\n",
    "    inputs.append(tokenizer(item['text'], max_length=300, truncation=True, \n",
    "                    padding=True, return_tensors='pt').to(device))\n",
    "    true_labels.append(item['label'])\n",
    "\n",
    "s_time = time.time()\n",
    "for i, input in enumerate(inputs):\n",
    "    logits   = model(**input)\n",
    "    label_id = np.argmax(logits[0].detach().cpu().numpy(), axis=1)[0]\n",
    "    pred_labels.append(label_id)\n",
    "print(\"avg time: \", (time.time() - s_time) * 1000 / len(raw_datasets))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 量化前精度评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0     0.9293    0.9214    0.9253     12500\n",
      "           1     0.9220    0.9299    0.9260     12500\n",
      "\n",
      "    accuracy                         0.9256     25000\n",
      "   macro avg     0.9257    0.9256    0.9256     25000\n",
      "weighted avg     0.9257    0.9256    0.9256     25000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import classification_report\n",
    "\n",
    "print(classification_report(true_labels, pred_labels, digits=4))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. PTQ动态量化\n",
    "\n",
    "## 2.1 设置量化后端"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.backends.quantized.supported_engines\n",
    "\n",
    "torch.backends.quantized.engine = 'x86'\n",
    "\n",
    "# 8-bit 量化\n",
    "cpu_device      = torch.device(\"cpu\")\n",
    "quantized_model = torch.quantization.quantize_dynamic(\n",
    "    model.to(cpu_device), {torch.nn.Linear}, dtype=torch.qint8\n",
    ").to(cpu_device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 量化性能评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "avg time:  0.21846478462219238\n"
     ]
    }
   ],
   "source": [
    "pred_labels =  [] \n",
    "inputs_cpu  = [input.to(cpu_device) for input in inputs]\n",
    "\n",
    "s_time = time.time()\n",
    "for i, input_cpu in enumerate(inputs_cpu[:100]):\n",
    "    logits   = quantized_model(**input_cpu)\n",
    "    label_id = np.argmax(logits[0].detach().cpu().numpy(), axis=1)[0]\n",
    "    pred_labels.append(label_id)\n",
    "print(\"avg time: \", (time.time() - s_time) * 1000 / len(raw_datasets))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3量化精度评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import classification_report\n",
    "\n",
    "# print(classification_report(true_labels, pred_labels, digits=4))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.4 模型大小对比"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size (MB):  433.312206\n",
      "Size (MB):  176.791466\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "def print_size_of_model(model):\n",
    "    torch.save(model.state_dict(), \"temp.p\")\n",
    "    print(\"Size (MB): \", os.path.getsize(\"temp.p\")/1e6)\n",
    "    os.remove(\"temp.p\")\n",
    "\n",
    "print_size_of_model(model)\n",
    "print_size_of_model(quantized_model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "总结： 不支持CUDA后端，毫无用处啊！"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
